/******************************************************************/
/*                                                                */
/*  Winpooch : Windows Watchdog                                   */
/*  Copyright (C) 2004-2005  Benoit Blanchon                      */
/*                                                                */
/*  This program is free software; you can redistribute it        */
/*  and/or modify it under the terms of the GNU General Public    */
/*  License as published by the Free Software Foundation; either  */
/*  version 2 of the License, or (at your option) any later       */
/*  version.                                                      */
/*                                                                */
/*  This program is distributed in the hope that it will be       */
/*  useful, but WITHOUT ANY WARRANTY; without even the implied    */
/*  warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR       */
/*  PURPOSE.  See the GNU General Public License for more         */
/*  details.                                                      */
/*                                                                */
/*  You should have received a copy of the GNU General Public     */
/*  License along with this program; if not, write to the Free    */
/*  Software Foundation, Inc.,                                    */
/*  675 Mass Ave, Cambridge, MA 02139, USA.                       */
/*                                                                */
/******************************************************************/


/******************************************************************/
/* Build configuration                                            */
/******************************************************************/

#define SPYDLL_BUILD_DLL

#define	TRACE_LEVEL	2

//#define SAVE_SPY_INFO


/******************************************************************/
/* Includes                                                       */
/******************************************************************/

// module's interface
#include "Hooks.h"

// standard headers
#include <windows.h>
#include <stdio.h>
#include <tchar.h>
#include <psapi.h>

// project's headers
#include "Assert.h"
#include "Disasm.h"
#include "HookNet.h"
#include "HookSys.h"
#include "HookReg.h"
#include "HookFile.h"
#include "SpyDll.h"
#include "SpyDllInt.h"
#include "Trace.h"


/******************************************************************/
/* Internal constants                                             */
/******************************************************************/

#define MAX_STUB_SIZE	20

#define MAX_BASE_NAME	32


/******************************************************************/
/* Internal data types                                            */
/******************************************************************/

typedef struct {
  BOOL		bEnabled ;
  UINT		nModuleId ;
  LPCSTR	szProcNameA ;
  PROC		pfnSpy ;
  BOOL		bHooked ;
  PROC		pfnTarget ;
  BYTE		pStub[MAX_STUB_SIZE] ;
  int		nStubSize ;    
} HOOKSFUNCTIONSTRUCT ;

typedef struct {
  LPCTSTR	szName ;
  HMODULE	hModule ;
} HOOKSMODULESTRUCT ;


/******************************************************************/
/* Internal data                                                  */
/******************************************************************/

static HOOKSMODULESTRUCT g_aModules[_HOOKS_MODULE_COUNT] = 
  {
    { TEXT("kernel32.dll") },
    { TEXT("advapi32.dll") },
    { TEXT("ws2_32.dll") }
  } ;
    

static HOOKSFUNCTIONSTRUCT g_aFunctions[_HOOKS_FUNCTION_COUNT] =
  {
    { TRUE,	HOOKS_KERNEL32_DLL,	"CreateProcessA",	(PROC)Spy_CreateProcessA	},
    { TRUE,	HOOKS_KERNEL32_DLL,	"CreateProcessW",	(PROC)Spy_CreateProcessW	},
    { TRUE,	HOOKS_ADVAPI32_DLL,	"RegCreateKeyExA",	(PROC)Hook_RegCreateKeyExA	},
    { TRUE,	HOOKS_ADVAPI32_DLL,	"RegCreateKeyExW",	(PROC)Hook_RegCreateKeyExW	},
    { TRUE,	HOOKS_ADVAPI32_DLL,	"RegCloseKey",		(PROC)Hook_RegCloseKey		},
    { TRUE,	HOOKS_ADVAPI32_DLL,	"RegOpenKeyExA",	(PROC)Hook_RegOpenKeyExA	},
    { TRUE,	HOOKS_ADVAPI32_DLL,	"RegOpenKeyExW",	(PROC)Hook_RegOpenKeyExW	},
    { TRUE,	HOOKS_ADVAPI32_DLL,	"RegSetValueExA",	(PROC)Hook_RegSetValueExA	},
    { TRUE,	HOOKS_ADVAPI32_DLL,	"RegSetValueExW",	(PROC)Hook_RegSetValueExW	},
    { TRUE,	HOOKS_KERNEL32_DLL,	"CreateFileW",		(PROC)Hook_CreateFileW		},
    { TRUE,	HOOKS_KERNEL32_DLL,	"DeleteFileW",		(PROC)Hook_DeleteFileW		},
    { TRUE,	HOOKS_KERNEL32_DLL,	"MoveFileWithProgressW",(PROC)Hook_MoveFileWithProgressW },
    { TRUE,	HOOKS_WS2_32_DLL,	"connect",		(PROC)Spy_connect		},
    { TRUE,	HOOKS_WS2_32_DLL,	"listen",		(PROC)Spy_listen		}
  } ;

VOID	* g_aStubs[_HOOKS_FUNCTION_COUNT] ;

/******************************************************************/

BOOL Hooks_Init ()
{
  int	i ;

#ifdef SAVE_SPY_INFO
  TCHAR szDir[64] ;

  // create dir for log
  wsprintf (szDir, TEXT("c:\\spytest\\%u"), GetCurrentProcessId()) ;
  CreateDirectory (TEXT("c:\\spytest"), NULL) ;
  CreateDirectory (szDir, NULL) ;
#endif

  for( i=0 ; i<_HOOKS_FUNCTION_COUNT ; i++ )
    g_aStubs[i] = g_aFunctions[i].pStub ;
  
  return HookReg_Init () ;
}

BOOL Hooks_Uninit ()
{
  
  return HookReg_Uninit () ;
}
/*
LPCTSTR Hooks_GetDllName (SPYFUNC id) 
{
  return aDescriptors[id].szDll ;
}

LPCSTR Hooks_GetProcName (SPYFUNC id) 
{
  return aDescriptors[id].szProc ;
}

PROC Hooks_GetProcAddress (SPYFUNC id)
{
  return aDescriptors[id].pfnSpy ;
}

VOID Hooks_SetStubAddress (SPYFUNC id, PROC pfnStub) 
{
  aDescriptors[id].pfnStub = pfnStub ;
}

BOOL	Hooks_IsEnabled		(SPYFUNC id) 
{
  return aDescriptors[id].bEnabled ;
}
*/



/******************************************************************/
/* Exported function                                              */
/******************************************************************/

PROC Hooks_GetStubAddress (HOOKS_FUNCTION id) 
{
  HOOKSFUNCTIONSTRUCT	*p ;

  ASSERT (id<_HOOKS_FUNCTION_COUNT) ;

  // get function descriptor
  p = &g_aFunctions[id] ;

  ASSERT (p->bHooked) ;

  return (PROC)p->pStub ;
}


/******************************************************************/
/* Exported function                                              */
/******************************************************************/

VOID* Hooks_GetStubAddForJmp (HOOKS_FUNCTION id) 
{
  ASSERT (id<_HOOKS_FUNCTION_COUNT) ;

  return &g_aStubs[id] ;
}


/******************************************************************/
/* Exported function                                              */
/******************************************************************/

BOOL Hooks_InstallIntoMappedModules	() 
{
  UINT nModuleId ;

  for( nModuleId=0 ; nModuleId<_HOOKS_MODULE_COUNT ; nModuleId++)
    Hooks_InstallIntoModule (nModuleId) ;

  return TRUE ;
}


/******************************************************************/
/* Exported function                                              */
/******************************************************************/

BOOL Hooks_InstallIntoModule (HOOKS_MODULE nModuleId) 
{
  UINT	nFunctionId ;

  ASSERT (nModuleId<_HOOKS_MODULE_COUNT) ;
 
  // already hooked ?
  if( g_aModules[nModuleId].hModule )
    return TRUE ;

  // get module handle
  g_aModules[nModuleId].hModule = GetModuleHandle(g_aModules[nModuleId].szName) ;
  
  // module mapped ?
  if( ! g_aModules[nModuleId].hModule ) 
    return FALSE ;
  
  DbgPrint (TEXT("Installing hooks in %s...\n"), 
	    g_aModules[nModuleId].szName) ;

  // hooks each function of that module
  for( nFunctionId=0 ; nFunctionId<_HOOKS_FUNCTION_COUNT ; nFunctionId++ )
    if( g_aFunctions[nFunctionId].nModuleId == nModuleId )
      if( g_aFunctions[nFunctionId].bEnabled )
	Hooks_InstallIntoFunction (nFunctionId) ;

  return TRUE ;
}


/******************************************************************/
/* Exported function                                              */
/******************************************************************/

BOOL Hooks_UninstallFromModule (HOOKS_MODULE nModuleId) 
{
  UINT	nFunctionId ;

  // really hooked ?
  if( ! g_aModules[nModuleId].hModule )
    return TRUE ;

  DbgPrint (TEXT("Removing hooks from %s...\n"), 
	    g_aModules[nModuleId].szName) ;

  // unhooks each function of that module
  for( nFunctionId=0 ; nFunctionId<_HOOKS_FUNCTION_COUNT ; nFunctionId++ )
    if( g_aFunctions[nFunctionId].nModuleId == nModuleId )
      if( g_aFunctions[nFunctionId].bHooked )
	 Hooks_UninstallFromFunction (nFunctionId) ;

  // remove module handle
  g_aModules[nModuleId].hModule = NULL ;

  return TRUE ;
}


/******************************************************************/
/* Exported function                                              */
/******************************************************************/

BOOL Hooks_DeleteFromModule (HOOKS_MODULE nModuleId) 
{
  UINT	nFunctionId ;

  // really hooked ?
  if( ! g_aModules[nModuleId].hModule )
    return TRUE ;

  DbgPrint (TEXT("Deleting hooks from %s...\n"), 
	    g_aModules[nModuleId].szName) ;

  // delete each function of that module
  for( nFunctionId=0 ; nFunctionId<_HOOKS_FUNCTION_COUNT ; nFunctionId++ )
    if( g_aFunctions[nFunctionId].nModuleId == nModuleId )
      if( g_aFunctions[nFunctionId].bHooked )
	 Hooks_DeleteFromFunction (nFunctionId) ;

  // remove module handle
  g_aModules[nModuleId].hModule = NULL ;

  return TRUE ;
}


/******************************************************************/
/* Exported function                                              */
/******************************************************************/

BOOL	Hooks_InstallIntoFunction (HOOKS_FUNCTION nFunctionId) 
{
  HOOKSFUNCTIONSTRUCT	*p ;
  MEMAREA		aForbidden[1] ;	
  MEMORY_BASIC_INFORMATION	mbiTarget;
  DWORD				dwDummy ;
  BOOL			bSuccess ;
  DWORD			dwJmpDst ;
  int			nHeadSize, nInstSize, i ;
#ifdef SAVE_SPY_INFO
  FILE	* fp ;
  CHAR szFile[64] ;
#endif

  ASSERT (nFunctionId<_HOOKS_FUNCTION_COUNT) ;

  TRACE_INFO (TEXT("nFunctionId = %u\n"), nFunctionId) ;

  // get function descriptor
  p = &g_aFunctions[nFunctionId] ;

  TRACE_INFO (TEXT("szProcName = %hs\n"), p->szProcNameA) ;

  // already hooked ?
  if( p->bHooked ) {
    TRACE_WARNING (TEXT("Function %u (%hs) already hooked\n"),
		   nFunctionId, p->szProcNameA) ;
    return FALSE ;
  }
     
  // get target address
  ASSERT (p->nModuleId<_HOOKS_MODULE_COUNT) ;
  ASSERT (g_aModules[p->nModuleId].hModule!=NULL) ;
  p->pfnTarget = GetProcAddress (g_aModules[p->nModuleId].hModule, 
				 p->szProcNameA) ;
  
  // failed ?
  if( ! p->pfnTarget ) {
    TRACE_ERROR (TEXT("Failed to get address of %hs\n"),
		 p->szProcNameA) ;
    return FALSE ;
  }
  
  // defined thread forbidden areas
  aForbidden[0].pBase	= p->pfnTarget ;
  aForbidden[0].nSize	= MAX_STUB_SIZE ;

  // supsend other threads
  SpyInt_SuspendOtherThreads (aForbidden, 1) ;

  //
  // Remove protection 
  // ----------------- 
  
  // get target protection
  VirtualQuery (p->pfnTarget, &mbiTarget, sizeof(MEMORY_BASIC_INFORMATION)) ;
  
  // change target protection
  bSuccess = VirtualProtect (mbiTarget.BaseAddress, mbiTarget.RegionSize,
			     PAGE_EXECUTE_READWRITE, &dwDummy) ; 
  if( ! bSuccess ) {
    TRACE_ERROR (TEXT("VirtualProtect failed (error=%u)\n"),
		 GetLastError()) ;
    goto failure ;
  }

#ifdef SAVE_SPY_INFO

  // save target head
  sprintf (szFile, "c:\\spytest\\%u\\%s-origin.hex",
	   GetCurrentProcessId(), p->szProcNameA) ;
  fp = fopen (szFile, "wb") ;
  fwrite (p->pfnTarget, MAX_STUB_SIZE, 1, fp) ;
  fclose (fp) ;

#endif

  //
  // Create stub
  // -----------
  
  // mesure head size
  nHeadSize = 0 ;
  while( nHeadSize < 5 ) 
    {
      nInstSize = Dasm_GetInstructionSize (((BYTE*)p->pfnTarget)+nHeadSize) ;
      if( ! nInstSize ) {
	TRACE_ERROR (TEXT("GetInstructionSize failed (id=%u, target=%hs)\n"), 
		     nFunctionId, p->szProcNameA) ;
	goto failure ;	
      }
      
      nHeadSize += nInstSize ;
      
      if( nHeadSize >= MAX_STUB_SIZE-5 ) {
	TRACE_ERROR (TEXT("Head too big\n")) ;
	goto failure ;
      }
      
    }
  
  // copy head into stub
  if( ! Dasm_MoveProgram (p->pStub, MAX_STUB_SIZE, (BYTE*)p->pfnTarget, nHeadSize) ) 
    {
      TRACE_ERROR (TEXT("MoveProgram failed\n")) ;
      goto failure ;
    }

  // add jmp at the end of the stub
  p->pStub[nHeadSize] = 0xE9 ; // JMP rel32
  dwJmpDst = (DWORD)p->pfnTarget - (DWORD)p->pStub - nHeadSize ;
  memcpy (p->pStub+nHeadSize+1, &dwJmpDst, 4) ;  
  
  // set size of stub
  p->nStubSize = nHeadSize + 5 ; // 5==sizeof(JMP rel32)

  // set stub as executable
  bSuccess = VirtualProtect (p->pStub, MAX_STUB_SIZE, 
			     PAGE_EXECUTE_READWRITE, &dwDummy) ; 
  if( ! bSuccess ) {
    TRACE_ERROR (TEXT("VirtualProtect failed (error=%u)\n"),
		 GetLastError()) ;
    goto failure ;
  }

#ifdef SAVE_SPY_INFO

  // save stub
  sprintf (szFile, "c:\\spytest\\%u\\%s-stub.hex", 
	   GetCurrentProcessId(), p->szProcNameA) ;
  fp = fopen (szFile, "wb") ;
  fwrite (p->pStub, p->nStubSize, 1, fp) ;
  fclose (fp) ;

#endif

  //
  // Patch target
  // ------------

  // overwrite head with a jmp to spy function
  ((BYTE*)p->pfnTarget)[0] = 0xE9 ; // JMP rel32
  dwJmpDst = (DWORD)p->pfnSpy - (DWORD)p->pfnTarget - 5 ; // 5==sizeof(JMP rel32)
  memcpy (((BYTE*)p->pfnTarget)+1, &dwJmpDst, 4) ;

  // complete head with NOPs
  for( i=5 ; i<nHeadSize ; i++ )
    ((BYTE*)p->pfnTarget)[i] = 0x90 ; // NOP
 
  p->bHooked = TRUE ;
  
#ifdef SAVE_SPY_INFO

  // save patched version of the target
  sprintf (szFile, "c:\\spytest\\%u\\%s-patched.hex",
	   GetCurrentProcessId(), p->szProcNameA) ;
  fp = fopen (szFile, "wb") ;
  fwrite (p->pfnTarget, MAX_STUB_SIZE, 1, fp) ;
  fclose (fp) ; 

  // save info about target
  sprintf (szFile, "c:\\spytest\\%u\\%s-info.txt", 
	   GetCurrentProcessId(), p->szProcNameA) ;
  fp = fopen (szFile, "wt") ;
  fprintf (fp, "Target = 0x%lX\n", (DWORD)p->pfnTarget) ;
  fprintf (fp, "Stub = 0x%lX\n", (DWORD)p->pStub) ;
  fprintf (fp, "Spy = 0x%lX\n", (DWORD)p->szProcNameA) ;
  fprintf (fp, "HeadSize = %d\n", nHeadSize) ;
  fclose (fp) ; 

#endif
  
  //
  // Restore protection
  // ------------------

  // restore target protection
  VirtualProtect (mbiTarget.BaseAddress, mbiTarget.RegionSize, 
		  mbiTarget.Protect, &dwDummy) ;  

  // resume other threads
  SpyInt_ResumeOtherThreads () ;
  
  return TRUE ;
  
 failure:

  p->bHooked = FALSE ;
    
  SpyInt_ResumeOtherThreads () ;
  
  return FALSE ;
}



/******************************************************************/
/* Exported function                                              */
/******************************************************************/

BOOL	Hooks_UninstallFromFunction (HOOKS_FUNCTION nFunctionId) 
{
  HOOKSFUNCTIONSTRUCT	*p ;
  MEMAREA		aForbidden[2] ;	
  DWORD			dwDummy ;
  MEMORY_BASIC_INFORMATION	mbi;

  ASSERT (nFunctionId<_HOOKS_FUNCTION_COUNT) ;

  TRACE_INFO (TEXT("nFunctionId = %u\n"), nFunctionId) ;

  // get function descriptor
  p = &g_aFunctions[nFunctionId] ;

  ASSERT (g_aModules[p->nModuleId].hModule!=NULL) ;

  // not hooked ?
  if( ! p->bHooked ) {
    TRACE_WARNING (TEXT("Function %u (%ls) not hooked\n"),
		   nFunctionId, p->szProcNameA) ;
    return FALSE ;
  }
  
  // defined thread forbidden areas
  aForbidden[0].pBase	= p->pfnTarget ;
  aForbidden[0].nSize	= MAX_STUB_SIZE ;
  aForbidden[1].pBase	= p->pStub ;
  aForbidden[1].nSize	= MAX_STUB_SIZE ;

  // supsend other threads
  SpyInt_SuspendOtherThreads (aForbidden, 2) ;

  // get target protection
  VirtualQuery (p->pfnTarget, &mbi, sizeof(MEMORY_BASIC_INFORMATION)) ;
  
  // change target protection
  VirtualProtect (mbi.BaseAddress, mbi.RegionSize, PAGE_EXECUTE_READWRITE, &mbi.Protect) ; 
  
  // copy head from stub
  memcpy (p->pfnTarget, p->pStub, p->nStubSize-5) ;
  
  // restore target protection
  VirtualProtect (mbi.BaseAddress, mbi.RegionSize, mbi.Protect, &dwDummy) ;  
  
  // resume other threads
  SpyInt_ResumeOtherThreads () ;

  p->bHooked = FALSE ;
  
  return TRUE ;
}


/******************************************************************/
/* Exported function                                              */
/******************************************************************/

BOOL	Hooks_DeleteFromFunction (HOOKS_FUNCTION nFunctionId) 
{
  HOOKSFUNCTIONSTRUCT	*p ;

  ASSERT (nFunctionId<_HOOKS_FUNCTION_COUNT) ;

  TRACE_INFO (TEXT("nFunctionId = %u\n"), nFunctionId) ;

  // get function descriptor
  p = &g_aFunctions[nFunctionId] ;

  ASSERT (g_aModules[p->nModuleId].hModule!=NULL) ;

  // not hooked ?
  if( ! p->bHooked ) {
    TRACE_WARNING (TEXT("Function %u (%ls) not hooked\n"),
		   nFunctionId, p->szProcNameA) ;
    return FALSE ;
  }

  p->bHooked = FALSE ;
  
  return TRUE ;
}


/******************************************************************/
/* Exported function                                              */
/******************************************************************/

BOOL Hooks_UninstallAll () 
{
  UINT	nModuleId ;
  UINT	nFunctionId ;

  TRACE ;

  // unhooks each module
  for( nModuleId=0 ; nModuleId<_HOOKS_MODULE_COUNT ; nModuleId++ )
    {
      if( g_aModules[nModuleId].hModule!=NULL )
	{
	  // unhooks each function
	  for( nFunctionId=0 ; nFunctionId<_HOOKS_FUNCTION_COUNT ; nFunctionId++ )
	    if( g_aFunctions[nFunctionId].nModuleId == nModuleId )
	      if( g_aFunctions[nFunctionId].bHooked )
		Hooks_UninstallFromFunction (nFunctionId) ;

	  g_aModules[nModuleId].hModule = NULL ;
	}
    }
     
  return TRUE ;
}
